package com.liuhengjia.springaiapply.api;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

/**
 * 模型聊天 API（控制层）
 */
@Slf4j
@RestController
@RequiredArgsConstructor
public class ChatAPI {
    private final ChatClient chatClient;
    private final OllamaChatModel ollamaChatModel;

    /**
     * 普通聊天（含RAG）
     *
     * @param message 用户输入
     * @return 响应
     */
    @GetMapping("/ai/chat")
    public String chat(@RequestParam(value = "message") String message) {
        log.error("/ai/chat: message: {}", message);
        return chatClient.prompt().user(message).call().content();
    }

    /**
     * 函数调用
     *
     * @param message 用户输入
     * @return 响应
     */
    @GetMapping("/ai/func")
    public String ragChat(@RequestParam(value = "message") String message) {
        return chatClient.prompt()
                .user(message)
                .functions("askForLeave")
                .call()
                .content();
    }

    /**
     * 多模态
     *
     * @param message 问题
     * @param file    图片
     * @return 响应
     */
    @PostMapping("/ai/model")
    public String model(@RequestParam(value = "message") String message,
                        @RequestParam(value = "file") MultipartFile file) {
        UserMessage userMessage = new UserMessage(
                message,
                List.of(new Media(MimeTypeUtils.IMAGE_JPEG, file.getResource()))
        );
        return ollamaChatModel.call(new Prompt(
                List.of(userMessage),
                ChatOptions.builder()
                        .model(OllamaModel.LLAVA.getName())
                        .build())
        ).getResult().getOutput().getText();
    }
}
